﻿using System;
using System.Collections.Generic;
using System.IdentityModel.Policy;
using System.Net;
using System.Runtime.Serialization;
using System.ServiceModel.Channels;
using System.ServiceModel.Security;
using System.ServiceModel.Web;
using System.Text;
using System.Xml;
using System.Xml.Linq;
using Microsoft.IdentityModel.Claims;
using Microsoft.IdentityModel.Tokens;
using Microsoft.ServiceModel.Web;

namespace Microsoft.IdentityModel.OAuth
{
    public class OAuthWrapSecurityChannel : RequestInterceptor
    {
        FederatedServiceCredentials wifCredentials;
        public OAuthWrapSecurityChannel()
            : base(true)
        {
        }

        public override void Initialize(System.ServiceModel.Channels.BindingContext context)
        {
            wifCredentials =  context.BindingParameters.Find<FederatedServiceCredentials>();
            if (wifCredentials == null)
                throw new InvalidOperationException("WIF is not configured on this service.");
        }
        public override void ProcessRequest(ref System.ServiceModel.Channels.RequestContext requestContext)
        {
            try
            {
                if (requestContext.RequestMessage.Properties.Via!= null &&
                    string.CompareOrdinal(requestContext.RequestMessage.Properties.Via.AbsolutePath.TrimStart('/'), "clientaccesspolicy.xml")==0)
                    return;

                var httpRequest = requestContext.RequestMessage.Properties[HttpRequestMessageProperty.Name] as HttpRequestMessageProperty;

                if (httpRequest == null)
                    throw new WebFaultException(HttpStatusCode.BadRequest);

                var authorizationHeader = httpRequest.Headers[HttpRequestHeader.Authorization];
                var accessToken = SWTParser.ExtractAndDecodeAccessToken(authorizationHeader);

                if (string.IsNullOrEmpty(accessToken))
                    throw new WebFaultException(HttpStatusCode.Unauthorized);

                using (XmlReader reader = WrapInsideBinarySecurityToken(accessToken))
                {
                    var handlersCollection = wifCredentials.SecurityTokenHandlerCollectionManager[SecurityTokenHandlerCollectionManager.Usage.Default];
                    if (handlersCollection.CanReadToken(reader))
                    {
                          //Read and validate token. This require SWT WIF handler to be plugged in
                        var token = handlersCollection.ReadToken(reader);
                        var identities = handlersCollection.ValidateToken(token);

                        //Allow the host to perform any claims-transformation...
                        var principal = wifCredentials.ClaimsAuthenticationManager.Authenticate(httpRequest.Method, new ClaimsPrincipal(identities));

                        var identityCollection = (principal != null) ? principal.Identities : new ClaimsIdentityCollection();
                        var authorizationPolicies = new List<IAuthorizationPolicy> { new AuthorizationPolicy(identityCollection) }.AsReadOnly();

                        //populate WCF SecurityContext.
                        var security = SecurityMessageProperty.GetOrCreate(requestContext.RequestMessage);
                        security.ServiceSecurityContext = new System.ServiceModel.ServiceSecurityContext(authorizationPolicies);
                    }
                }
            }
            catch (WebFaultException fault)
            {
                EndRequest(ref requestContext, requestContext.RequestMessage, fault.StatusCode);
            }
            catch (WebFaultException<HttpError> typedFault)
            {
                var reply = Message.CreateMessage(MessageVersion.None, "", new HttpErrorBodyWriter(typedFault.Detail));
                EndRequest(ref requestContext, reply, typedFault.StatusCode);
            }
        }

        class HttpErrorBodyWriter : BodyWriter
        {
            DataContractSerializer dcs;
            HttpError errorObject;
            public HttpErrorBodyWriter(HttpError error)
                : base(true)
            {
                this.errorObject = error;
                this.dcs = new DataContractSerializer(error.GetType());
            }
            protected override void OnWriteBodyContents(XmlDictionaryWriter writer)
            {
                dcs.WriteObject(writer, errorObject);
            }
        }

        private XmlReader WrapInsideBinarySecurityToken(string accessToken)
        {
            var root = new XElement(
                XNamespace.Get("http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-wssecurity-secext-1.0.xsd").GetName("BinarySecurityToken"),
                new XAttribute("ValueType", "http://schemas.xmlsoap.org/ws/2009/11/swt-token-profile-1.0"),
                new XAttribute("EncodingType", "http://docs.oasis-open.org/wss/2004/01/oasis-200401-wss-soap-message-security-1.0#Base64Binary"),
                Convert.ToBase64String(Encoding.Default.GetBytes(accessToken)));

            return root.CreateReader();
        }

        private static HttpResponseMessageProperty EnsureHttpResposne(Message message)
        {
            HttpResponseMessageProperty httpResponse;
            object rmp;
            message.Properties.TryGetValue(HttpResponseMessageProperty.Name, out rmp);
            httpResponse = rmp as HttpResponseMessageProperty;
            if (httpResponse == null)
            {
                httpResponse = new HttpResponseMessageProperty();
                message.Properties.Add(HttpResponseMessageProperty.Name, httpResponse);
            }
            return httpResponse;
        }

        private static void EndRequest(ref RequestContext requestContext, Message reply, HttpStatusCode code)
        {
            if (reply == null)
                reply = requestContext.RequestMessage;

            var httpResponse = EnsureHttpResposne(reply);

            httpResponse.StatusCode = code;
            requestContext.Reply(reply);
            requestContext = null;
        }
    }
}
